#include "data_set.h"

DataSet::DataSet() {
    _activityLog = "\n";
    _totalFreqs = new AlleleFrequencies();
    addToLog( "DataSet initiated." );
    _treeWidgetItem = new QTreeWidgetItem();
    _treeWidgetItem->setText(0,QObject::tr("Genotypes"));
}

DataSet::DataSet(QList<Individual*> inds) {
    _activityLog = "\n";
    _totalFreqs = new AlleleFrequencies();
    addToLog( QString("DataSet initiated."));
    addIndividuals(inds);

}



void DataSet::addToLog(QString msg) {
    _activityLog += QString("[%1]: %5\n").arg( QTime::currentTime().toString("hh:mm:ss.zzz")).arg(msg);
}


void DataSet::removedMonomorphicLoci() {
    QStringList loci = _individuals[0]->locusNames();
    QStringList fixedLoci;

    foreach(QString locus, loci ){
        if( _totalFreqs->getLocusFrequency(locus)->isFixed() ) {
            foreach(Individual *ind, _individuals )
                ind->removeLocus( locus);
            fixedLoci.append(locus);
        }
    }
    addToLog( QString("Removed Fixed Loci: %1").arg( fixedLoci.join(", ")));
}

void DataSet::addIndividual(Individual *ind){
    _individuals.append(ind);
    _totalFreqs->addIndividual( ind );
}

void DataSet::addIndividuals(QList<Individual*>inds) {
    foreach( Individual* ind, inds )
        _individuals.append( ind );
    addToLog( QString("Added %1 individuals to data set").arg(inds.count()));
}

Individual* DataSet::getIndividual( int idx ) {
    if( idx < _individuals.count() )
        return _individuals[idx];
    else
        return new Individual();
}

QStringList DataSet::getStrataKeys() {
    return _individuals[0]->strataNames();
}

int DataSet::popSize(QString key) {
    return (_populations.keys().contains(key)) ? _populations[key].count() : 0;
}

void DataSet::partitionByStratum(QString key) {
    resetPartitions();
    foreach(Individual *ind, _individuals ){
        QString val = ind->getStrata( key );
        if( !_populations.keys().contains( val ) ){
            QList<Individual*> inds;
            _popFreqs[val] = new AlleleFrequencies();
            inds.append( ind );
            _populations[val] = inds;
        }
        else
            _populations[val].append( ind );
        _popFreqs[val]->addIndividual( ind );
    }
    addToLog( QString("Partitioned data on: %1").arg(key));
}


void DataSet::resetPartitions() {
    _populations.clear();
    _popFreqs.clear();
    addToLog( "Reset Partitions" );
}

AlleleFrequencies* DataSet::getTotalFrequencies() {
    return _totalFreqs;
}

AlleleFrequencies* DataSet::getStratumFrequencies( QString key ) {
    if( !_popFreqs.keys().contains(key))
        return new AlleleFrequencies;
    else
        return _popFreqs[key];
}
